/**
 * @file MeshProcessor.cpp
 * @brief 网格处理器的实现文件
 * 
 * 该文件实现了网格数据的读取、处理和插值计算的核心功能。
 * 包含了并行计算优化和内存管理优化。
 */

#include "MeshProcessor.h"
#include "FileIO.h"
#include "Interpolator.h"
#include <iostream>
#include <iomanip>
#include <thread>
#include <algorithm>
#include <array>
#ifdef _OPENMP
#include <omp.h>  // OpenMP并行计算支持
#endif
#ifdef _WIN32
#include <windows.h>
#endif

// 初始化静态成员变量，用于存储K近邻点的数量
size_t MeshProcessor::K_NEAREST_POINTS = 0;

/**
 * @brief 网格处理器构造函数
 * @param grid_a 源网格文件路径
 * @param grid_b 目标网格文件路径
 * @param props_a 源网格属性文件路径
 * @param output 输出文件路径
 * 
 * 初始化网格处理器并设置OpenMP线程数
 */
MeshProcessor::MeshProcessor(const std::string &grid_a,
                           const std::string &grid_b,
                           const std::string &props_a,
                           const std::string &output)
    : grid_a_file(grid_a),
      grid_b_file(grid_b),
      props_a_file(props_a),
      output_file(output)
{
    #ifdef _OPENMP
    // 设置OpenMP使用所有可用的处理器核心
    omp_set_num_threads(omp_get_num_procs());
    #endif
}

/**
 * @brief 设置K近邻点数量
 * @param k 需要使用的近邻点数量
 * 
 * 设置插值计算时使用的最近邻点数量
 */
void MeshProcessor::setKNearestPoints(size_t k)
{
    K_NEAREST_POINTS = k;
}

/**
 * @brief 验证网格数据的有效性
 * @param size 网格尺寸
 * @param coord 坐标数据
 * @param zcorn 深度角点数据
 * @throw std::runtime_error 当数据大小不符合要求时抛出异常
 * 
 * 检查COORD和ZCORN数据的大小是否符合网格尺寸要求
 */
void MeshProcessor::validateGridData(const GridSize &size,
                                   const std::vector<std::vector<double>> &coord,
                                   const std::vector<double> &zcorn)
{
    size_t expected_coord_size = (size_t)(size.Nx + 1) * (size.Ny + 1);
    size_t expected_zcorn_size = (size_t)8 * size.Nx * size.Ny * size.Nz;

    if (coord.size() != expected_coord_size)
    {
        throw std::runtime_error("COORD数据大小不正确");
    }

    if (zcorn.size() != expected_zcorn_size)
    {
        throw std::runtime_error("ZCORN数据大小不正确");
    }
}

/**
 * @brief 计算网格单元中心点坐标
 * @param size 网格尺寸
 * @param coord 坐标数据
 * @param zcorn 深度角点数据
 * @return 返回所有网格单元中心点的坐标
 * 
 * 该函数执行以下步骤:
 * 1. 重组ZCORN数据为3D数组便于访问
 * 2. 并行计算每个网格单元的8个顶点坐标
 * 3. 计算每个单元的中心点坐标(8个顶点的平均值)
 */
std::vector<Point3D> MeshProcessor::calculateCellCenters(const GridSize &size,
                                                       const std::vector<std::vector<double>> &coord,
                                                       const std::vector<double> &zcorn)
{
    std::vector<Point3D> centers(size.Nx * size.Ny * size.Nz);

    // 重塑ZCORN数据为3D数组
    std::vector<std::vector<std::vector<double>>> zcorn_3d(2 * size.Nz,
        std::vector<std::vector<double>>(2 * size.Ny,
            std::vector<double>(2 * size.Nx)));

    int idx = 0;
    for (int k = 0; k < 2 * size.Nz; ++k)
        for (int j = 0; j < 2 * size.Ny; ++j)
            for (int i = 0; i < 2 * size.Nx; ++i)
                zcorn_3d[k][j][i] = zcorn[idx++];

    #pragma omp parallel for collapse(3)
    for (int k = 0; k < size.Nz; ++k)
    {
        for (int j = 0; j < size.Ny; ++j)
        {
            for (int i = 0; i < size.Nx; ++i)
            {
                int cell_idx = k * size.Nx * size.Ny + j * size.Nx + i;
                int k_offset = 2 * k;
                int j_offset = 2 * j;
                int i_offset = 2 * i;

                // 获取单元八个顶点的坐标
                std::vector<Point3D> vertices(8);
                
                // 底面四个顶点
                vertices[0] = {coord[j * (size.Nx + 1) + i][0],
                             coord[j * (size.Nx + 1) + i][1],
                             zcorn_3d[k_offset][j_offset][i_offset]};
                             
                vertices[1] = {coord[j * (size.Nx + 1) + (i + 1)][0],
                             coord[j * (size.Nx + 1) + (i + 1)][1],
                             zcorn_3d[k_offset][j_offset][i_offset + 1]};
                             
                vertices[2] = {coord[(j + 1) * (size.Nx + 1) + i][0],
                             coord[(j + 1) * (size.Nx + 1) + i][1],
                             zcorn_3d[k_offset][j_offset + 1][i_offset]};
                             
                vertices[3] = {coord[(j + 1) * (size.Nx + 1) + (i + 1)][0],
                             coord[(j + 1) * (size.Nx + 1) + (i + 1)][1],
                             zcorn_3d[k_offset][j_offset + 1][i_offset + 1]};

                // 顶面四个顶点
                vertices[4] = {coord[j * (size.Nx + 1) + i][0],
                             coord[j * (size.Nx + 1) + i][1],
                             zcorn_3d[k_offset + 1][j_offset][i_offset]};
                             
                vertices[5] = {coord[j * (size.Nx + 1) + (i + 1)][0],
                             coord[j * (size.Nx + 1) + (i + 1)][1],
                             zcorn_3d[k_offset + 1][j_offset][i_offset + 1]};
                             
                vertices[6] = {coord[(j + 1) * (size.Nx + 1) + i][0],
                             coord[(j + 1) * (size.Nx + 1) + i][1],
                             zcorn_3d[k_offset + 1][j_offset + 1][i_offset]};
                             
                vertices[7] = {coord[(j + 1) * (size.Nx + 1) + (i + 1)][0],
                             coord[(j + 1) * (size.Nx + 1) + (i + 1)][1],
                             zcorn_3d[k_offset + 1][j_offset + 1][i_offset + 1]};

                // 中心点
                Point3D &center = centers[cell_idx];
                for (const auto &v : vertices)
                {
                    center.x += v.x;
                    center.y += v.y;
                    center.z += v.z;
                }
                center.x /= 8.0;
                center.y /= 8.0;
                center.z /= 8.0;
            }
        }
    }

    return centers;
}

/**
 * @brief 计算目标点的插值结果
 * @param target 目标点坐标
 * @param neighbors K个最近邻点及其距离
 * @return 返回插值计算的杨氏模量和泊松比
 * 
 * 使用反距离加权(IDW)方法进行插值:
 * 1. 如果存在在距离为0的点，直接使用该点的值
 * 2. 使用SIMD指令优化权重计算
 * 3. 计算加权平均值作为插值结果
 */
std::pair<double, double> MeshProcessor::calculateInterpolation(
    const Point3D &target,
    const std::vector<std::pair<double, KDTree::KDNode *>> &neighbors)
{
    const size_t n = neighbors.size();
    if (n == 0) return {0.0, 0.0};

    // 如果距离为0，直接返回该点的值
    for (const auto &neighbor : neighbors)
    {
        if (neighbor.first < 1e-10)
        {
            return {neighbor.second->young_value, neighbor.second->poisson_value};
        }
    }

    // 使用反距离加权插值
    double sum_weights = 0.0;
    double sum_young = 0.0;
    double sum_poisson = 0.0;

    // 每4个点为一组进行SIMD计算
    const size_t simd_size = (n / 4) * 4;
    for (size_t i = 0; i < simd_size; i += 4)
    {
        // 加载4个距离值
        __m256d dist = _mm256_set_pd(
            neighbors[i+3].first,
            neighbors[i+2].first,
            neighbors[i+1].first,
            neighbors[i].first
        );

        // 计算权重 (1/d^2)
        __m256d weight = _mm256_div_pd(_mm256_set1_pd(1.0),
                                      _mm256_mul_pd(dist, dist));

        // 加载属性值
        __m256d young = _mm256_set_pd(
            neighbors[i+3].second->young_value,
            neighbors[i+2].second->young_value,
            neighbors[i+1].second->young_value,
            neighbors[i].second->young_value
        );

        __m256d poisson = _mm256_set_pd(
            neighbors[i+3].second->poisson_value,
            neighbors[i+2].second->poisson_value,
            neighbors[i+1].second->poisson_value,
            neighbors[i].second->poisson_value
        );

        // 计算加权和
        __m256d weighted_young = _mm256_mul_pd(weight, young);
        __m256d weighted_poisson = _mm256_mul_pd(weight, poisson);

        // 累加结果
        double w[4], y[4], p[4];
        _mm256_store_pd(w, weight);
        _mm256_store_pd(y, weighted_young);
        _mm256_store_pd(p, weighted_poisson);

        for (int j = 0; j < 4; ++j)
        {
            sum_weights += w[j];
            sum_young += y[j];
            sum_poisson += p[j];
        }
    }

    // 处理剩余的点
    for (size_t i = simd_size; i < n; ++i)
    {
        double dist = neighbors[i].first;
        double weight = 1.0 / (dist * dist);
        sum_weights += weight;
        sum_young += weight * neighbors[i].second->young_value;
        sum_poisson += weight * neighbors[i].second->poisson_value;
    }

    // 避免除零
    if (sum_weights < 1e-10)
    {
        return {0.0, 0.0};
    }

    return {sum_young / sum_weights, sum_poisson / sum_weights};
}

/**
 * @brief 执行插值计算
 * @throw std::runtime_error 当插值计算失败时抛出异常
 * @throw std::bad_alloc 当内存分配失败时抛出异常
 * 
 * 处理流程:
 * 1. 构建源网格点的KD树
 * 2. 使用线程本地存储优化并行处理
 * 3. 对目标点进行分块处理提高缓存命中
 * 4. 执行K近邻搜索和插值计算
 * 5. 收集并存储结果
 */
void MeshProcessor::performInterpolation()
{
    // 预分配并对齐内存
    const int points_size = static_cast<int>(centers_a.size());
    std::vector<std::pair<Point3D, std::pair<double, double>>> points_data;
    points_data.reserve(points_size);

    // 使用线程本地存储优化
    #pragma omp parallel
    {
        alignas(32) std::vector<std::pair<Point3D, std::pair<double, double>>> local_points;
        const int thread_id = omp_get_thread_num();
        const int num_threads = omp_get_num_threads();
        const int chunk_size = (points_size + num_threads - 1) / num_threads;
        const int start = thread_id * chunk_size;
        const int end = std::min(start + chunk_size, points_size);
        
        local_points.reserve(chunk_size);

        for (int i = start; i < end; ++i)
        {
            if (i < static_cast<int>(props_a.youngmod.size()) && 
                i < static_cast<int>(props_a.poissonr.size()))
            {
                local_points.emplace_back(centers_a[i],
                    std::make_pair(props_a.youngmod[i], props_a.poissonr[i]));
            }
        }

        #pragma omp critical
        {
            points_data.insert(points_data.end(), 
                std::make_move_iterator(local_points.begin()),
                std::make_move_iterator(local_points.end()));
        }
    }

    kdtree.build(points_data);
    points_data.clear();
    points_data.shrink_to_fit();

    // 初始化结果数组
    const int result_size = static_cast<int>(centers_b.size());
    results.resize(result_size);

    // 用缓存友好的分块处理
    constexpr int BLOCK_SIZE = 64;
    const int num_blocks = (result_size + BLOCK_SIZE - 1) / BLOCK_SIZE;

    #pragma omp parallel for schedule(guided)
    for (int block = 0; block < num_blocks; ++block)
    {
        const int start = block * BLOCK_SIZE;
        const int end = std::min(start + BLOCK_SIZE, result_size);

        alignas(32) std::vector<std::pair<double, KDTree::KDNode*>> nearest;
        nearest.reserve(K_NEAREST_POINTS);

        for (int i = start; i < end; ++i)
        {
            nearest.clear();
            kdtree.findKNearest(centers_b[i], K_NEAREST_POINTS, nearest);
            auto interp_result = calculateInterpolation(centers_b[i], nearest);
            results[i] = Result{i, centers_b[i],
                              interp_result.first, interp_result.second};
        }
    }
}

/**
 * @brief 处理网格属性插值的主函数
 * @throw std::exception 当处理过程中出现错误时抛出异常
 * 
 * 执行完整的处理流程:
 * 1. 预分配内存并设置大页面支持
 * 2. 读取源网格数据和属性
 * 3. 计算源网格中心点
 * 4. 读取目标网格数据
 * 5. 计算目标网格中心点
 * 6. 执行插值计算
 * 7. 保存结果
 */
void MeshProcessor::process()
{
    try
    {
        // 预分配内存
        centers_a.reserve(size_a.Nx * size_a.Ny * size_a.Nz);
        centers_b.reserve(size_b.Nx * size_b.Ny * size_b.Nz);
        results.reserve(size_b.Nx * size_b.Ny * size_b.Nz);
        
        // 设置大页面支持
        #ifdef _WIN32
        HANDLE process = GetCurrentProcess();
        SIZE_T min_ws_size = 0, max_ws_size = 0;
        GetProcessWorkingSetSize(process, &min_ws_size, &max_ws_size);
        SetProcessWorkingSetSize(process, 
            min_ws_size + 256 * 1024 * 1024,
            max_ws_size + 512 * 1024 * 1024);
        #endif

        // 1. 读取源网格数据
        FileIO::readGridFile(grid_a_file, size_a, coord_a, zcorn_a);
        FileIO::readPropertiesFile(props_a_file, props_a);

        // 2. 计算源网格中心点
        centers_a = calculateCellCenters(size_a, coord_a, zcorn_a);

        // 3. 读取目标网格数据
        FileIO::readGridFile(grid_b_file, size_b, coord_b, zcorn_b);

        // 4. 计算目标网格中心点
        centers_b = calculateCellCenters(size_b, coord_b, zcorn_b);

        // 5. 执行插值计算
        performInterpolation();

        // 6. 保存结果
        FileIO::saveResults(output_file, results);
    }
    catch (const std::exception &e)
    {
        throw;
    }
}